//
//  堆.swift
//  Swift-LeetCode
//
//  Created by 卢悦明 on 2024/2/25.
//

import UIKit
/**
 二叉堆本质上一个“完全二叉树” 获取最大/最小的数，复杂度O（1）删除和添加都是logn
 大顶堆：子节点小于父节点
 小顶堆：子节点大于父节点
 二叉堆是通过数组实现的
 */
class BinaryHeap<E: Equatable & Comparable> {
    
    enum BinaryHeap_Type {
        case max_binary_heap // 大顶堆
        case mix_binary_heap // 小顶对
    }
    private var count = 0
    private let defaultValue: E
    var elements: [E]
    private let binaryHeapType: BinaryHeap_Type
    
    
    init(_ length: Int = 10, defaultValue: E, _ type: BinaryHeap_Type = .max_binary_heap) {
        self.defaultValue = defaultValue
        self.binaryHeapType = type
        elements = Array<E>(repeating: defaultValue, count: length)
    }
    
    /// 批量建堆（自上而下的上滤）
    /// - Parameters:
    ///   - array: 无序的数组
    ///   - defaultValue: 默认值
    init(_ array: [E], defaultValue: E,  _ type: BinaryHeap_Type = .max_binary_heap) {
        self.defaultValue = defaultValue
        self.binaryHeapType = type
        elements = Array(array)
        count = elements.count
        for index in 1..<count {
            siftUp(index)
        }
    }
    
    /// 批量建堆（自下而上的下滤）
    /// - Parameters:
    ///   - array: 无序的数组
    ///   - defaultValue: 默认值
    init(array: [E], defaultValue: E, _ type: BinaryHeap_Type = .max_binary_heap) {
        self.defaultValue = defaultValue
        self.binaryHeapType = type
        elements = Array(array)
        count = elements.count
        let leaf = (count/2) - 1
        for index in (0...leaf).reversed() {
            siftDown(index)
        }
    }
    
    func size() -> Int {
        count
    }
    func isEmpty() -> Bool {
        count == 0
    }
    func clear() {
        elements = elements.map { item in
            if item != defaultValue {
                return defaultValue
            } else {
                return item
            }
        }
        count = 0
    }
    func add(_ elem: E) {
        creatMaxArray(count + 1)
        elements[count] = elem
        count += 1
        /** 上滤
         1、添加的elem大于父节点，两个交换位置
         2、小于等于父节点或者父节点没有父节点，不移动退出循环
         */
        siftUp(count - 1)

        
    }
//    获取堆顶元素
    func get() -> E? {
        elements.first
    }
    // 删除堆顶元素
    @discardableResult func remove() -> E? {
        if count == 0 {
            return nil
        }
        let temp = elements.first
        elements[0] = elements[count - 1]
        elements[count - 1] = defaultValue
        count -= 1
        /** 下滤
         如果有子节点，子节点中较大的节点大于当前节点，当前节点和子节点中较大的节点交换
         2、小于等于父节点或者父节点没有父节点，不移动退出循环
         */
        siftDown(0)
        return temp
    }
    // 删除堆顶元素，并添加一个新元素
    @discardableResult func replace(_ elem: E) -> E? {
        var root:E?
        if count == 0 {
            elements[0] = elem
            count += 1
        } else {
            root = elements[0]
            elements[0] = elem
            siftDown(0)
        }
        return root
    }
}

extension BinaryHeap {
//    index位置的元素上滤
    private func siftUp(_ index: Int) {
        if index <= 0 { return }
        var newIndex = index
        let elemet = elements[newIndex]
        while newIndex > 0 {
            let parentIndex = (newIndex - 1) / 2;
            let parentElem = elements[parentIndex]
            if (compare(a: parentElem, b: elemet) >= 0) {
                break
            }
            elements[newIndex] = parentElem
            newIndex = parentIndex
        }
        elements[newIndex] = elemet
    }
//    下虑
    private func siftDown(_ index: Int) {
        let element = elements[index]
        let harl = count >> 1 // 非叶子节点的索引
        var index = index
        // 非叶子节点
        while index < harl {
            // 获取左子节点
            var childIndex = (index * 2) + 1
            var child = elements[childIndex]
            // 获取右子节点
            let rightIndex = childIndex + 1
            // 右节点存在
            if rightIndex < count,
                compare(a: elements[rightIndex], b: child) > 0 {
                child = elements[rightIndex]
                childIndex = rightIndex
            }
            if (compare(a: element, b: child) >= 0)  {
                break
            }
            // 交换
            elements[index] = child
            index = childIndex
        }
        elements[index] = element
    }
//    元素比较
    private func compare(a: E, b: E) -> Int {
        if binaryHeapType == .max_binary_heap {
            if a == b {
                return 0
            } else if a > b {
                return 1
            } else {
                return -1
            }
        } else {
            if a == b {
                return 0
            } else if a < b {
                return 1
            } else {
                return -1
            }
        }
    }
//    数组扩容
    private func creatMaxArray(_ macCount: Int) {
        let length = elements.count
        if length > macCount {
            return
        }
        let newLength = length + (length >> 1)
        var newArray = Array<E>(repeating: defaultValue, count: newLength)
        for (index, item) in elements.enumerated() {
            newArray[index] = item
        }
        elements = newArray
    }
    func printHeap() {
        if count == 0 {
            print("[]")
            return
        }
        var leval = 0
        var levalCount = 1
        var array = [E]()
        print("节点总数**********", count)
        for index in 0..<count {
            let item = elements[index]
            if levalCount == 0 {
                leval += 1
                levalCount = 1 << leval
                print(array)
                array.removeAll()
                array.append(item)
            } else {
                array.append(item)
            }
            levalCount -= 1
            if index == count-1 {
                print(array)
            }
        }
    }
}
